from pymilvus import connections, Collection
from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection
from sentence_transformers import SentenceTransformer
from .mongo import jokes_col
from config import Config

# -------------------------------
# Milvus 连接
# -------------------------------
MILVUS_CONN = Config.MILVUS_CONN  # 连接名
connections.connect(MILVUS_CONN, host=Config.MILVUS_HOST, port=Config.MILVUS_PORT)

# -------------------------------
# Collection 定义
# -------------------------------
COLLECTION_NAME = Config.MILVUS_COLLECTION
fields = [
    FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=50, is_primary=True, auto_id=False),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=Config.EMBED_DIM),
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=500)
]
schema = CollectionSchema(fields, description="AI Joker 笑话集合（使用MongoDB _id作为主键）")

# 删除旧集合（可选）
if COLLECTION_NAME in connections.list_connections():
    Collection(COLLECTION_NAME, using=MILVUS_CONN).drop()

# 创建新集合
col = Collection(COLLECTION_NAME, schema, using=MILVUS_CONN)

# -------------------------------
# 加载向量模型
# -------------------------------
model = SentenceTransformer(Config.EMBED_MODEL)

# -------------------------------
# 批量从 MongoDB 获取数据
# -------------------------------
BATCH_SIZE = Config.IMPORT_BATCH_SIZE

def fetch_mongo_data(batch_size=BATCH_SIZE):
    """按批次获取MongoDB文档"""
    cursor = jokes_col.find({}, {"_id": 1, "content": 1})
    batch_ids, batch_texts = [], []

    for doc in cursor:
        batch_ids.append(str(doc["_id"]))
        batch_texts.append(doc["content"])

        if len(batch_ids) >= batch_size:
            yield batch_ids, batch_texts
            batch_ids, batch_texts = [], []

    if batch_ids:
        yield batch_ids, batch_texts

# -------------------------------
# 插入 Milvus
# -------------------------------
total_count = 0
for ids, texts in fetch_mongo_data():
    embeddings = model.encode(texts, convert_to_numpy=True).tolist()
    col.insert([ids, embeddings, texts])
    total_count += len(ids)

# -------------------------------
# 创建索引
# -------------------------------
col.create_index(
    "embedding",
    {"index_type": Config.MILVUS_INDEX_TYPE, "metric_type": Config.MILVUS_METRIC_TYPE, "params": {"nlist": Config.MILVUS_NLIST}}
)

# 加载集合到内存
col.load()

print(f"✅ 已成功导入 {total_count} 条笑话数据到 Milvus，使用MongoDB的_id作为主键。")